package tests;

import java.io.FileWriter;
import java.io.IOException;
import java.io.PrintWriter;
import java.util.Arrays;
import java.util.HashSet;
import java.util.Random;

import stdp_rules.LogSTDP;
import utils.StateCamera;
import utils.StateCamera.CameraType;
import utils.SynapseRandomizer;
import utils.math_utils.Coordinate;
import utils.math_utils.MathUtils;
import utils.math_utils.RandStream;
import utils.math_utils.RandStream.Rand;
import Jama.Matrix;
import components.LIFNeuron;
import components.Layer;
import components.Network;
import components.Synapse;
import connectors.Connect;
import connectors.GridConnect;

public class XORTest {
	
	public static final int numXORDist = 15;
	public static final int trainEpochs = 5000;
	public static final int testEpochs = 1000;
	public static final double timeStep = 1; //ms
	public static final double XORisi = 15; //ms
	public static StateCamera camera;
	public static int numInputs = 400;
	public static int numRes = 1024;
	public static LIFNeuron[] inputs;
	public static LIFNeuron[] res;
	public static LIFNeuron[] res1;
	public static LIFNeuron[] res2;
	public static LIFNeuron[] res3;
	public static LIFNeuron[] res4;
	public static Synapse [][] rr;
	public static Synapse [][] ir;
	public static boolean [] trainingIn;
	public static boolean [] testingIn;
	public static double [][] teachTargs;
	public static double [][] testTargs;
	public static LogSTDP stdp;
	public static GridConnect gCon;
	
	public static void initializeNetwork() {
		
		camera = new StateCamera(CameraType.Bin,
				timeStep, XORisi);
		
		gCon = new GridConnect();
		
		inputs = initializeInputs(numInputs);

		res1 = initializeNeurons(256, 30, 15, 2.5,  0.3, camera);
		res2 = initializeNeurons(256, 30, 15, 2.5,  0.3, camera);
		res3 = initializeNeurons(256, 30, 15, 2.5,  0.3, camera);
		res4 = initializeNeurons(256, 30, 15, 2.5,  0.3, camera);
		
		res = MathUtils.concatArrays(res1, res2);
		res = MathUtils.concatArrays(res, res3);
		res = MathUtils.concatArrays(res, res4);
		rr = new Synapse [numRes][numRes];

		initializeSynVars(gCon);
			
		Network net = new Network();
		Layer<LIFNeuron> ins = new Layer<LIFNeuron>();
		ins.setNeurons(inputs);
		Layer<LIFNeuron> rsvr = new Layer<LIFNeuron>();
		rsvr.setNeurons(res);
		
		
		inputs = ins.getNeuronArray(new LIFNeuron[numInputs]);
		res = rsvr.getNeuronArray(new LIFNeuron [numRes]);
		
		int [] inDims = {20, 20, 1};
		int [] resDims = {8, 8, 16};
		net.cubicLayout(new Coordinate(0,0,0), inDims , inputs);
		net.cubicLayout(new Coordinate(6, 6, 2), resDims, res);

		
		gCon.stochasticConnect(res, res, 2.25);
//		gCon.stochasticConnect(res1, res1, 1.75);
//		
//		gCon.stochasticConnect(res1, res2, 2);
//		gCon.stochasticConnect(res2, res1, 2);
//		
//		gCon.stochasticConnect(res2, res2, 2.25);
//		
//		gCon.stochasticConnect(res2, res3, 2);
//		gCon.stochasticConnect(res3, res2, 2);
//		
//		gCon.stochasticConnect(res3, res3, 2.25);
//		
//		gCon.stochasticConnect(res3, res4, 2);
//		gCon.stochasticConnect(res4, res3, 2);
//		
//		gCon.stochasticConnect(res4, res4, 1.75);
		
		int y = 0;
		int i = 0;
		for(LIFNeuron n : res) {
			n.setIndex(i);
			i++;
		}
		i = 0;
		for(LIFNeuron n : res) {
			n.setIndex(y);
			for(Synapse s: n.getFanOut()) {
				rr[y][s.getTarg().getIndex()] = s;
				i++;
			}
			y++;
		}
		int [] degree = new int[numRes];
		for(int p = 0; p < numRes; p++) {
			degree[p] = res[p].getFanIn().size() + res[p].getFanOut().size();
		}
		System.out.println(Arrays.toString(degree));
		System.out.println(i);
		ir = gCon.stochasticConnect(inputs, res, 16);
		double lRate = 0.3 / Math.pow(10, 3);
		                // tcP,  tcD,  lRat,  nv,   alp,  bet, cp, wc 
		stdp = new LogSTDP(5,   10, lRate, 0.01f, 3.0f, 30f, 1f, 20.0f);
//		for(int q = 0; q < numRes; q++) {
//			res[q].transferSynapsesToArray();
//		}
//		for(int r = 0; r < numInputs; r++) {
//			inputs[r].transferSynapsesToArray();
//		}
	}	
	
	/**
	 * Generates the XOR test/train inputs/targets
	 */
	public static void generateXORPattern() {
		Random rand = new Random();
		trainingIn = new boolean[trainEpochs];
		teachTargs = new double[trainEpochs - (numXORDist+1)][numXORDist];
		Matrix inXOR = new Matrix(trainEpochs, 1);
		for (int i = 0; i < trainEpochs; i++) {
			trainingIn[i] = rand.nextBoolean();
			inXOR.set(i, 0, trainingIn[i] ? 1 : 0);
			if (i > numXORDist) {
				for (int j = 1; j <= numXORDist; j++) {
					if (i >= j) {
						teachTargs[i-(numXORDist+ 1)][j-1] = (trainingIn[i] != trainingIn[i-j]) ? 1 : 0;
					}
				}
			}
		}
		
		PrintWriter pw = null;
		try{
			FileWriter fw = new FileWriter("xorIns.dat");
			pw = new PrintWriter(fw);
			inXOR.print(pw, 4, 3);
		} catch (IOException e) {
			e.printStackTrace();
		} finally {
			if(pw != null)
				pw.close();
		}
		
		
		testingIn = new boolean[testEpochs];
		testTargs = new double[testEpochs - (numXORDist+1)][numXORDist];
		for (int i = 0; i < testEpochs; i++) {
			testingIn[i] = rand.nextBoolean();
			if (i > numXORDist) {
				for (int j = 1; j <= numXORDist; j++) {
					if (i >= j) {
						testTargs[i-(numXORDist+ 1)][j-1] = (testingIn[i] != testingIn[i-j]) ? 1 : 0;
					}
				}
			}
		}	
		
	}
	
	/**
	 * Run an XOR through time test.
	 * @param args
	 */
	public static void main(String [] args) {
		initializeNetwork();
		
		
	//***********************************************************************//
	//*************************CONSISTENCY CHECK****************************//
//		int counter = 0;
		Matrix RR = new Matrix(rr.length, rr[0].length);
		for (int i = 0, n = rr.length; i < n; i++) {	
			for (int j = 0, m = rr[i].length; j < m; j++) {
				if (rr[i][j] != null) {
					RR.set(i, j, rr[i][j].getW());
//					counter ++;
					if(Double.isInfinite(rr[i][j].getR())) {
						System.out.println("R Problem rr " + rr[i][j].getR());
					}
					if(rr[i][j].getD() < 0) {
						System.out.println("D Problem rr " + rr[i][j].getD());
					}
				}
				
			}	
		}
		try {
			FileWriter FW = new FileWriter("ReservoirMatBefore.dat");
			PrintWriter pw = new PrintWriter(FW);
			
			RR.print(pw, 4, 3);
			pw.close();
		} catch (IOException e) {
			e.printStackTrace();
		}
		
//		int nullCounter1 = 0;
//		int nullCounter2 = 0;
//		int synNCounter = 0;
//		for (LIFNeuron n : res) {
//			if(n.getFanIn() != null) {
//				nullCounter1++;
//				for(Synapse s: n.getFanOut()) {
//					if (s == null) {
//						synNCounter++;
//					}
//				}
//			}
//			if(n.getFanOut() != null) {
//				nullCounter2++;
//			}
//		}
		
		
//		System.out.println("RR # Non-Null FanIn " + nullCounter1);
//		System.out.println("RR # Non-Null FanOut " + nullCounter2);
//		System.out.println("RR # Null FanIn Syns " + synNCounter);
//		System.out.println("RR # of Syns Total: " + counter);
//		System.out.println("RR Sparity: " + (float) counter/(rr.length * rr[0].length));
	
//		counter = 0;
//		for (int i = 0, n = ir.length; i < n; i++) {	
//			for (int j = 0, m = ir[i].length; j < m; j++) {
//				if (ir[i][j] != null) {
//					counter ++;
//					if(Double.isInfinite(ir[i][j].getR())) {
//						System.out.println("R Problem ir " + ir[i][j].getR());
//					}
//					if(ir[i][j].getD() < 0) {
//						System.out.println("D Problem ir " + ir[i][j].getU());
//					}
//				}
//			}
//		}
//		nullCounter1 = 0;
//		nullCounter2 = 0;
//		for (LIFNeuron n : inputs) {
//			if(n.getFanIn() != null) {
//				nullCounter1++;
//			}
//			if(n.getFanOut() != null) {
//				nullCounter2++;
//			}
//		}
//		System.out.println("IR # Non-Null FanIn " + nullCounter1);
//		System.out.println("IR # Non-Null FanOut " + nullCounter2);
//		System.out.println("IR # of Syns Total: " + counter);
//		System.out.println("IR Sparity: " + (float) counter/(ir.length * ir[0].length));

		

		
	//***************************************************************//
	//*************************FIRST-PASS****************************//
		
		camera.readyFilm(trainEpochs * (int)XORisi, numRes);
		HashSet<Integer> inds = new HashSet<Integer>();
		for(int i = 0; i < numRes; i++) {
			inds.add(i);
		}

		int count = 0;
		
		generateXORPattern();
		
		System.out.println("First-Pass...");
		int perDone = 0;
		double [] fireTimes = new double [numInputs];
		
		for(int i = 0; i < trainEpochs * XORisi/2; i++) {
			if(i % XORisi == 0) {
				for(int j = 0; j < numInputs; j++) {
					if(trainingIn[(int)(i/XORisi)] != ((j % 2) == 0))
						fireTimes[j] += XORisi;
					}
			}
			
			if(i % ((trainEpochs*XORisi/2)/10) == 0) {
				System.out.print(perDone + "%  ");
				perDone += 10;
			}

			for (int j = 0; j < numInputs; j++) {			
				if(fireTimes[j] + inputs[j].getLastSpike()
						<= inputs[j].getTime()){
					inputs[j].spike();
					stdp.applySTDPRule(inputs[j].getFanOut(), inputs[j].getTime()-timeStep, false);
					fireTimes[j] = XORisi/5 + RandStream.nextPoissonInterval(1);
				}
				inputs[j].update(timeStep);
			}

			
			for (int q = 0; q < numRes; q++) {
				if(res[q].update(timeStep)){
					stdp.applySTDPRule(res[q].getFanOut(), res[q].getTime()-timeStep, false);
					stdp.applySTDPRule(res[q].getFanIn(), res[q].getTime()-timeStep, true);
					if(inds.contains(q)){
						inds.remove(q);
					}
					count++;
				}
			}
			//camera.directSnapshot();
		}
		//camera.printStatesToFile("fulOuts.dat");
		
		Matrix RS = new Matrix(rr.length, rr[0].length);
		int y = 0;
		for(LIFNeuron n : res) {
			for(Synapse s: n.getFanOut()) {
				RS.set(y, s.getTarg().getIndex(), s.getW());
			}
			y++;
		}
		PrintWriter PW = null;
		 try {
				FileWriter fw = new FileWriter("fullOuts.dat");
				PW = new PrintWriter(fw);
				RR.print(PW, 4, 3);
				FileWriter fw2 = new FileWriter("postSTDP.dat");
				PW = new PrintWriter(fw2);
				RS.print(PW, 4, 3);
				fw.close();
				PW.close();
						
			} catch(IOException e) {
				e.printStackTrace();
			} finally {
				if(PW != null)
					PW.close();
			}

		 
	//**************************************************************//
	//*************************TRAINING****************************//
		System.out.println("\nTraining...");
		perDone = 0;
		camera.readyFilm(trainEpochs, numRes);
		for(int i = 0; i < trainEpochs * XORisi; i++) {
			if(i % XORisi == 0) {
				for(int j = 0; j < numInputs; j++) {
					if(trainingIn[(int)(i/XORisi)] != (j % 2 == 0))
						fireTimes[j] += XORisi;
					}
			}

			for (int j = 0; j < numInputs; j++) {
				if(fireTimes[j] + inputs[j].getLastSpike()
						<= inputs[j].getTime()){
					inputs[j].spike();
					fireTimes[j] = XORisi/5 + RandStream.nextPoissonInterval(1);
				}
					inputs[j].update(timeStep);
			}		
			
			if(i % ((trainEpochs*XORisi)/10) == 0) {
				System.out.print(perDone + "%  ");
				perDone += 10;
			}
			
			for (int q = 0; q < numRes; q++) {
				if(res[q].update(timeStep)){
					if(inds.contains(q)){
						inds.remove(q);
					}
					count++;
				}
			}
			camera.collectFrame();
		}
		
		System.out.println("\nSolving...");
		 double [][] states = Arrays.copyOfRange(camera.getCapturedStates(), (numXORDist+ 1), camera.getCapturedStates().length);
		 Matrix S = new Matrix(states);
		 Matrix D = new Matrix(teachTargs);
		 Matrix wOut = S.solve(D);
		

		System.out.println(count);
		System.out.println("Dead: " + inds.size());
		for(Integer I : inds) {
			System.out.println("Dead " + I +": " + res[I].getMemPotential());		
		}
		
		
		
		
//	//**************************************************************//
//	//*************************TESTING******************************//
		System.out.println("Testing...");
		perDone = 0;
		camera.readyFilm(testEpochs, numRes);

		for(int i = 0; i < testEpochs * XORisi; i++) {
			if(i % XORisi == 0) {
				for(int j = 0; j < numInputs; j++) {
					if(testingIn[(int)(i/XORisi)] != (j % 2 == 0))
						fireTimes[j] += XORisi;
					}
			}
			
			if(i % ((testEpochs*XORisi)/10) == 0) {
				System.out.print(perDone + "%  ");
				perDone += 10;
			}

			for (int j = 0; j < numInputs; j++) {
				if(fireTimes[j] + inputs[j].getLastSpike()
						<= inputs[j].getTime()){
					inputs[j].spike();
					fireTimes[j] = XORisi/5 + RandStream.nextPoissonInterval(1);
				}
				inputs[j].update(timeStep);
			}

			//Update RESERVOIR
			for (int q = 0; q < numRes; q++) {
				if(res[q].update(timeStep)){
					if(inds.contains(q)){
						inds.remove(q);
					}
					count++;
				}
			}
			camera.collectFrame();
		}
		
		
		states = Arrays.copyOfRange(camera.getCapturedStates(), (numXORDist+ 1), camera.getCapturedStates().length);
		S = new Matrix(states);
		Matrix tOuts = S.times(wOut);
		for(int i = 0; i < tOuts.getRowDimension(); i++) {
			for(int j = 0; j < numXORDist; j++) {
				tOuts.set(i, j, tOuts.get(i,j) >= 0.5 ? 1.0 : 0.0);
			}
		}	
		Matrix dOuts = new Matrix(testTargs);
		Matrix x = (tOuts.minus(dOuts));
		x.arrayTimesEquals(x);
		System.out.println();
		
//		try {
//			FileWriter fw = new FileWriter("TestTargets.dat");
//			PrintWriter pw = new PrintWriter(fw);
//			dOuts.print(pw, 3, 4);
//			pw.close();
//			
//			FileWriter FW = new FileWriter("TestOuts.dat");
//			PrintWriter P = new PrintWriter(FW);
//			tOuts.print(P,  3, 4);
//			P.close();
//			
//			
//		}catch (IOException e) {
//			e.printStackTrace();
//		}
////		
//		//FINDING ERRORS***********************************************
//		
		double [] errors = new double[numXORDist];
		for (int i = 0; i < numXORDist; i++){
			double val = 0.0;
			for(int j = 0; j < testEpochs - (numXORDist+ 1); j++) {
				val = x.get(j, i) + val;
			}
			errors[i] = val/(double)testEpochs;
		}
		
		System.out.println(Arrays.toString(errors));	
		
//		//**************************************************************//
//		//*************************RECORDING****************************//
//		
		camera.printStatesToFile("testSerial.dat");	
//		try {
//			FileWriter fw = new FileWriter("testOuts.dat");
//			FileWriter fw2 = new FileWriter("testTargs.dat");
//			PrintWriter pw = new PrintWriter(fw);
//			PrintWriter pw2 = new PrintWriter(fw2);
//			for(int i = 0; i < tOuts.getRowDimension(); i++){
//				pw.println(Arrays.toString(tOuts.getArray()[i]));
//				pw2.println(Arrays.toString(testTargs[i]));
//			}
//			pw.close();
//			pw2.close();
//		}catch(IOException e) {
//			e.printStackTrace();
//		}
		RR = new Matrix(rr.length, rr[0].length);
		for (int i = 0, n = rr.length; i < n; i++) {	
			for (int j = 0, m = rr[i].length; j < m; j++) {
				if (rr[i][j] != null) {
					RR.set(i, j, rr[i][j].getW());
//					counter ++;
					if(Double.isInfinite(rr[i][j].getR())) {
						System.out.println("R Problem rr " + rr[i][j].getR());
					}
					if(rr[i][j].getD() < 0) {
						System.out.println("D Problem rr " + rr[i][j].getD());
					}
				}
				
			}	
		}
		try {
			FileWriter FW = new FileWriter("ReservoirMatAfter.dat");
			PrintWriter pw = new PrintWriter(FW);
			
			RR.print(pw, 4, 3);
			pw.close();
		} catch (IOException e) {
			e.printStackTrace();
		}
	
	}
	
	/** 
	 * Values from Maass (2002), D & F in mSecs
	 */
	public static void initializeSynVars (Connect gCon) {
		
		SynapseRandomizer eeSR = new SynapseRandomizer();
		eeSR.setU(new RandStream(Rand.GAUSS, 0.5, 0.25, 0.001,
				Double.MAX_VALUE, true));
		eeSR.setD(new RandStream(Rand.GAUSS, 1100, 550, 0.001,
				Double.MAX_VALUE, true));
		eeSR.setF(new RandStream(Rand.GAUSS, 50, 25, 0.001,
				Double.MAX_VALUE, true));
		
		
		SynapseRandomizer eiSR = new SynapseRandomizer();
		eiSR.setU(new RandStream(Rand.GAUSS, 0.05, 0.025, 0.001,
				Double.MAX_VALUE, true));
		eiSR.setD(new RandStream(Rand.GAUSS, 125, 62.5, 0.001,
				Double.MAX_VALUE, true));
		eiSR.setF(new RandStream(Rand.GAUSS, 120, 60, 0.001,
				Double.MAX_VALUE, true));
		
		
		SynapseRandomizer ieSR = new SynapseRandomizer();
		ieSR.setU(new RandStream(Rand.GAUSS, 0.25, 0.125, 0.001,
				Double.MAX_VALUE, true));
		ieSR.setD(new RandStream(Rand.GAUSS, 700, 350, 0.001,
				Double.MAX_VALUE, true));
		ieSR.setF(new RandStream(Rand.GAUSS, 20, 10, 0.001,
				Double.MAX_VALUE, true));
		
		
		SynapseRandomizer iiSR = new SynapseRandomizer();
		iiSR.setU(new RandStream(Rand.GAUSS, 0.32, 0.16, 0.0001,
				Double.MAX_VALUE, true));
		iiSR.setD(new RandStream(Rand.GAUSS, 144, 72, 0.001,
				Double.MAX_VALUE, true));
		iiSR.setF(new RandStream(Rand.GAUSS, 60, 30, 0.001,
				Double.MAX_VALUE, true));
		
		gCon.setEeSR(eeSR);
		gCon.setEiSR(eiSR);
		gCon.setIeSR(ieSR);
		gCon.setIiSR(iiSR);
		
	}
	
	public static LIFNeuron [] initializeNeurons(int numNeurons, int tc,
			double meanRate, double stdRate, double percentInhib, StateCamera camera) {
		LIFNeuron [] neurons = new LIFNeuron[numNeurons];	

		//double tc;
		for (int i = 0; i < numNeurons; i++) {
			//tc = RandStream.nextUniform(20, 40);
			if(Math.random() < percentInhib) {
				neurons[i] = new LIFNeuron(tc, false, i);
				neurons[i].setRefractory(2);
			} else {
				neurons[i] = new LIFNeuron(tc, true, i);
				neurons[i].setRefractory(3);
			}
			neurons[i].setPreferredISI(RandStream.nextGaussian(meanRate, stdRate, 5, meanRate*2, true));
			neurons[i].setBgCurrent(13.5);
			neurons[i].setHmP(true);
			if(camera != null) {
				neurons[i].setReporter(true);
				neurons[i].setRecorder(camera);
			}
		}
		
		initializeNeuronVars(neurons);
		return neurons;
	}
	
	public static LIFNeuron [] initializeInputs(int numInputs) {
		LIFNeuron [] neurons = new LIFNeuron[numInputs];	
		for (int i = 0; i < numInputs; i++) {
			neurons[i] = new LIFNeuron(30, true, i);
			neurons[i].setRefractory(0);
		}
		initializeNeuronVars(neurons);
		return neurons;
	}
	
	public static void initializeNeuronVars (LIFNeuron [] neurons) {
		
		for (int i = 0, n = neurons.length; i < n; i++) {
			neurons[i].setThreshold(15);
			neurons[i].setLastSpike(0);
			neurons[i].setTime(0);
			neurons[i].setReset(13.5);
			neurons[i].setMemResistance(1);
			neurons[i].setRestPotential(0);	
			neurons[i].setMemPotential(RandStream.nextUniform(13.5, 15.1));	
		}		
		
	}
}
